#!/usr/bin/env python3
"""
RNN翻译器测试脚本
"""

import os
import sys
import torch
import pytest

# 添加路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.translator import RNTranslator, EncoderRNN, DecoderRNN, AttentionDecoderRNN

class TestRNTranslator:
    """RNN翻译器测试类"""
    
    def setup_method(self):
        """测试设置"""
        self.translator = RNTranslator(max_length=20, hidden_size=64, num_layers=1)
        
        # 创建简单的词汇表
        self.en_sentences = ["hello world", "how are you"]
        self.zh_sentences = ["你好世界", "你好吗"]
    
    def test_preprocess_text_en(self):
        """测试英文文本预处理"""
        text = "Hello, World! How are you?"
        result = self.translator.preprocess_text(text, 'en')
        expected = ['hello', ',', 'world', '!', 'how', 'are', 'you', '?']
        assert result == expected
    
    def test_preprocess_text_zh(self):
        """测试中文文本预处理"""
        text = "你好，世界！今天怎么样？"
        result = self.translator.preprocess_text(text, 'zh')
        expected = ['你', '好', '，', '世', '界', '！', '今', '天', '怎', '么', '样', '？']
        assert result == expected
    
    def test_build_vocab(self):
        """测试词汇表构建"""
        self.translator.build_vocab(self.en_sentences, self.zh_sentences)
        
        # 检查特殊标记
        assert self.translator.en_word2idx['<PAD>'] == 0
        assert self.translator.en_word2idx['<SOS>'] == 1
        assert self.translator.en_word2idx['<EOS>'] == 2
        assert self.translator.en_word2idx['<UNK>'] == 3
        
        # 检查词汇表大小
        assert len(self.translator.en_word2idx) > 4
        assert len(self.translator.zh_word2idx) > 4
    
    def test_sentence_to_indices(self):
        """测试句子转索引"""
        self.translator.build_vocab(self.en_sentences, self.zh_sentences)
        
        indices = self.translator.sentence_to_indices("hello world", 'en')
        
        # 检查长度
        assert len(indices) == self.translator.max_length
        
        # 检查特殊标记
        assert indices[0] == self.translator.SOS_token
        assert self.translator.EOS_token in indices
    
    def test_indices_to_sentence(self):
        """测试索引转句子"""
        self.translator.build_vocab(self.en_sentences, self.zh_sentences)
        
        # 创建测试索引
        test_indices = [self.translator.SOS_token, 
                       self.translator.zh_word2idx.get('你', self.translator.UNK_token),
                       self.translator.zh_word2idx.get('好', self.translator.UNK_token),
                       self.translator.EOS_token]
        
        sentence = self.translator.indices_to_sentence(test_indices, 'zh')
        assert isinstance(sentence, str)
    
    def test_create_model(self):
        """测试模型创建"""
        self.translator.build_vocab(self.en_sentences, self.zh_sentences)
        self.translator.create_model(use_attention=False)
        
        assert self.translator.encoder is not None
        assert self.translator.decoder is not None
        
        # 检查模型设备
        assert next(self.translator.encoder.parameters()).device == self.translator.device
    
    def test_encoder_forward(self):
        """测试编码器前向传播"""
        self.translator.build_vocab(self.en_sentences, self.zh_sentences)
        self.translator.create_model(use_attention=False)
        
        # 创建测试输入
        test_indices = self.translator.sentence_to_indices("hello", 'en')
        input_tensor = torch.tensor([test_indices], dtype=torch.long).to(self.translator.device)
        
        # 前向传播
        outputs, hidden = self.translator.encoder(input_tensor)
        
        # 检查输出形状
        assert outputs.shape[0] == 1  # batch size
        assert outputs.shape[1] == self.translator.max_length  # sequence length
        assert outputs.shape[2] == self.translator.hidden_size  # hidden size
    
    def test_decoder_forward(self):
        """测试解码器前向传播"""
        self.translator.build_vocab(self.en_sentences, self.zh_sentences)
        self.translator.create_model(use_attention=False)
        
        # 创建测试输入
        test_indices = self.translator.sentence_to_indices("你", 'zh')
        input_tensor = torch.tensor([test_indices[:1]], dtype=torch.long).to(self.translator.device)
        
        # 创建隐藏状态
        hidden = self.translator.encoder.init_hidden(1, self.translator.device)
        
        # 前向传播
        output, new_hidden = self.translator.decoder(input_tensor, hidden)
        
        # 检查输出形状
        assert output.shape[0] == 1  # batch size
        assert output.shape[1] == len(self.translator.zh_word2idx)  # vocabulary size

class TestEncoderRNN:
    """编码器测试类"""
    
    def test_encoder_initialization(self):
        """测试编码器初始化"""
        encoder = EncoderRNN(input_size=100, hidden_size=64, num_layers=2)
        
        assert encoder.hidden_size == 64
        assert encoder.num_layers == 2
        assert encoder.embedding is not None
        assert encoder.gru is not None
    
    def test_encoder_forward_shape(self):
        """测试编码器前向传播形状"""
        encoder = EncoderRNN(input_size=100, hidden_size=64, num_layers=2)
        
        # 创建测试输入
        input_seq = torch.randint(0, 100, (2, 10))  # batch_size=2, seq_len=10
        
        outputs, hidden = encoder(input_seq)
        
        # 检查输出形状
        assert outputs.shape == (2, 10, 64)  # (batch, seq_len, hidden_size)
        assert hidden.shape == (2, 2, 64)   # (num_layers, batch, hidden_size)

class TestDecoderRNN:
    """解码器测试类"""
    
    def test_decoder_initialization(self):
        """测试解码器初始化"""
        decoder = DecoderRNN(hidden_size=64, output_size=50, num_layers=2)
        
        assert decoder.hidden_size == 64
        assert decoder.output_size == 50
        assert decoder.num_layers == 2
        assert decoder.embedding is not None
        assert decoder.gru is not None
        assert decoder.out is not None
    
    def test_decoder_forward_shape(self):
        """测试解码器前向传播形状"""
        decoder = DecoderRNN(hidden_size=64, output_size=50, num_layers=2)
        
        # 创建测试输入
        input_seq = torch.randint(0, 50, (2, 1))  # batch_size=2, seq_len=1
        hidden = torch.zeros(2, 2, 64)  # (num_layers, batch, hidden_size)
        
        output, new_hidden = decoder(input_seq, hidden)
        
        # 检查输出形状
        assert output.shape == (2, 50)  # (batch, output_size)
        assert new_hidden.shape == (2, 2, 64)  # (num_layers, batch, hidden_size)

class TestAttentionDecoderRNN:
    """注意力解码器测试类"""
    
    def test_attention_decoder_initialization(self):
        """测试注意力解码器初始化"""
        decoder = AttentionDecoderRNN(hidden_size=64, output_size=50, num_layers=2)
        
        assert decoder.hidden_size == 64
        assert decoder.output_size == 50
        assert decoder.num_layers == 2
        assert decoder.attention is not None
    
    def test_attention_decoder_forward_shape(self):
        """测试注意力解码器前向传播形状"""
        decoder = AttentionDecoderRNN(hidden_size=64, output_size=50, num_layers=2)
        
        # 创建测试输入
        input_seq = torch.randint(0, 50, (2, 1))  # batch_size=2, seq_len=1
        hidden = torch.zeros(2, 2, 64)  # (num_layers, batch, hidden_size)
        encoder_outputs = torch.randn(2, 10, 64)  # (batch, seq_len, hidden_size)
        
        output, new_hidden, attention_weights = decoder(input_seq, hidden, encoder_outputs)
        
        # 检查输出形状
        assert output.shape == (2, 50)  # (batch, output_size)
        assert new_hidden.shape == (2, 2, 64)  # (num_layers, batch, hidden_size)
        assert attention_weights.shape == (2, 10, 1)  # (batch, encoder_seq_len, 1)

def test_translator_integration():
    """集成测试"""
    translator = RNTranslator(max_length=15, hidden_size=32, num_layers=1)
    
    # 简单训练测试
    en_sentences = ["hello", "world"]
    zh_sentences = ["你好", "世界"]
    
    # 构建词汇表
    translator.build_vocab(en_sentences, zh_sentences)
    
    # 创建模型
    translator.create_model(use_attention=False)
    
    # 测试翻译（虽然模型未训练，但应该能运行）
    try:
        result = translator.translate("hello")
        assert isinstance(result, str)
    except Exception as e:
        # 模型未训练时翻译会失败，这是正常的
        assert "模型未训练" in str(e) or "未找到" in str(e)

if __name__ == "__main__":
    # 运行测试
    pytest.main([__file__, "-v"])